import torch
import torchvision.datasets
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

train_data = torchvision.datasets.CIFAR10("./dataset", train=True, download=False,
                                          transform=torchvision.transforms.ToTensor())
train_dataloader = DataLoader(train_data, batch_size=1, shuffle=True, num_workers=0, drop_last=False)

test_data = torchvision.datasets.CIFAR10("./dataset", train=False, download=False,
                                         transform=torchvision.transforms.ToTensor())
test_dataloader = DataLoader(test_data, batch_size=1, shuffle=True, num_workers=0, drop_last=False)


class MyMod(nn.Module):
    def __init__(self):
        super(MyMod, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=32, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(in_features=1024, out_features=64),
            nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        output = self.model(x)
        return output


mymod = MyMod()
print(mymod)
input = torch.ones((64, 3, 32, 32))
output = mymod(input)
writer = SummaryWriter("../logs")
writer.add_graph(mymod, input)
writer.close()
print("test output: ", output.shape)

for data in train_dataloader:
    imgs, targets = data
    print(targets)
    print(imgs.shape)
    '''
    # 打印每层的结果，方便对比哪层出现问题
    x, targets = data
    print(x.shape)
    for i in range(len(mymod.model)):
        x = mymod.model[i](x)
        print(i, ' layer shape：', x.shape)
    '''
    output = mymod(imgs)
    print(output.shape)
    print(output)
    break
